import tool.TreeNode;

import java.util.*;

/**
 * 给出二叉树的根节点root，树上每个节点都有一个不同的值。
 * 如果节点值在to_delete中出现，我们就把该节点从树上删去，最后得到一个森林（一些不相交的树构成的集合）。
 * 返回森林中的每棵树。你可以按任意顺序组织答案。
 * <p>
 * 示例：
 * 输入：root = [1,2,3,4,5,6,7], to_delete = [3,5]
 * 输出：[[1,2,null,4],[6],[7]]
 * <p>
 * 提示：
 * 树中的节点数最大为1000。
 * 每个节点都有一个介于1 到1000之间的值，且各不相同。
 * to_delete.length <= 1000
 * to_delete 包含一些从1 到1000、各不相同的值。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode-cn.com/problems/delete-nodes-and-return-forest
 * 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
 */
public class Q1110m {

    private List<TreeNode> ans;

    private Set<Integer> dels;

    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        ans = new ArrayList<>();
        dels = new HashSet<>(to_delete.length << 1);
        for (int td : to_delete) {
            dels.add(td);
        }
        if (!dels.contains(root.val)) {
            ans.add(root);
        }
        postOrderDfs(root);
        return ans;
    }

    /**
     * 后续遍历
     */
    private void postOrderDfs(TreeNode root) {
        if (root.left != null) {
            postOrderDfs(root.left);
            if (dels.contains(root.left.val)) {
                root.left = null;
            }
        }
        if (root.right != null) {
            postOrderDfs(root.right);
            if (dels.contains(root.right.val)) {
                root.right = null;
            }
        }

        if (dels.contains(root.val)) {
            if (root.left != null) {
                ans.add(root.left);
            }
            if (root.right != null) {
                ans.add(root.right);
            }
            root.left = root.right = null;
        }
    }

    /**
     * 7ms 太慢了
     *
     * @param root
     * @param to_delete
     * @return
     */
    public List<TreeNode> delNodes1(TreeNode root, int[] to_delete) {
        List<TreeNode> ans = new ArrayList<>();
        if (root == null) return ans;
        Set<Integer> dels = new HashSet<>(to_delete.length << 1);
        for (int td : to_delete) {
            dels.add(td);
        }
        Map<TreeNode, TreeNode> map = new HashMap<>();
        map.put(root, new TreeNode(0));
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (queue.size() > 0) {
            root = queue.poll();
            if (root.left != null) {
                map.put(root.left, root);
                queue.offer(root.left);
            }
            if (root.right != null) {
                map.put(root.right, root);
                queue.offer(root.right);
            }
        }
        for (Map.Entry<TreeNode, TreeNode> entry : map.entrySet()) {
            if (dels.contains(entry.getKey().val) || dels.contains(entry.getValue().val)) {
                if (entry.getKey() == entry.getValue().left) {
                    entry.getValue().left = null;
                } else {
                    entry.getValue().right = null;
                }
            }
        }
        for (Map.Entry<TreeNode, TreeNode> entry : map.entrySet()) {
            if (!dels.contains(entry.getKey().val) && (entry.getValue().left == null && entry.getValue().right == null)) {
                ans.add(entry.getKey());
            }
        }
        return ans;
    }

}
